#include "websocket_client.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "esp_log.h"
#include "esp_camera.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include "lwip/sockets.h"
#include "lwip/netdb.h"
#include "esp_system.h"

static const char *TAG = "websocket_client";

// WebSocket客户端相关变量
static int websocket_sock = -1;
static bool websocket_connected = false;
static TaskHandle_t websocket_task_handle = NULL;
static TaskHandle_t websocket_frame_task_handle = NULL;
static bool websocket_task_running = false;
static bool websocket_frame_task_running = false;
static SemaphoreHandle_t websocket_mutex = NULL;  // 添加互斥锁


// 函数声明
static void start_websocket_frame_task(void);
static void stop_websocket_frame_task(void);
static void websocket_frame_task(void *pvParameters);
static void websocket_task(void *pvParameters);
static void generate_websocket_key(char *key, size_t key_size);
static bool websocket_handshake(int sock);
static bool websocket_send_binary(int sock, const uint8_t *data, size_t len);

// WebSocket密钥生成
static void generate_websocket_key(char *key, size_t key_size) {
    const char charset[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    for (int i = 0; i < 22; i++) {
        key[i] = charset[esp_random() % (sizeof(charset) - 1)];
    }
    key[22] = '=';
    key[23] = '=';
    key[24] = '\0';
}

// WebSocket握手
static bool websocket_handshake(int sock) {
    char key[25];
    generate_websocket_key(key, sizeof(key));
    
    char request[512];
    snprintf(request, sizeof(request),
        "GET %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: %s\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n", 
        WEBSOCKET_ENDPOINT, WEBSOCKET_SERVER_IP, WEBSOCKET_SERVER_PORT, key);
    
    int sent = send(sock, request, strlen(request), 0);
    if (sent < 0) {
        ESP_LOGE(TAG, "WebSocket发送握手请求失败");
        return false;
    }
    
    char response[1024];
    int received = recv(sock, response, sizeof(response) - 1, 0);
    if (received < 0) {
        ESP_LOGE(TAG, "WebSocket接收握手响应失败");
        return false;
    }
    response[received] = '\0';
    
    if (strstr(response, "101 Switching Protocols") && 
        strstr(response, "Upgrade: websocket")) {
        ESP_LOGI(TAG, "WebSocket握手成功");
        return true;
    }
    
    ESP_LOGE(TAG, "WebSocket握手失败: %s", response);
    return false;
}

// 发送WebSocket二进制数据帧（优化版本）
static bool websocket_send_binary(int sock, const uint8_t *data, size_t len) {
    if (len > 65535) {
        ESP_LOGE(TAG, "数据太大，不能发送");
        return false;
    }
    
    uint8_t frame[4];
    int frame_len = 0;
    
    // FIN=1, RSV=000, Opcode=0010 (binary)
    frame[0] = 0x82;
    
    if (len < 126) {
        frame[1] = 0x80 | (uint8_t)len;  // MASK=1, payload length
        frame_len = 2;
    } else {
        frame[1] = 0x80 | 126;  // MASK=1, payload length = 126
        frame[2] = (len >> 8) & 0xFF;
        frame[3] = len & 0xFF;
        frame_len = 4;
    }
    
    // 发送帧头
    if (send(sock, frame, frame_len, 0) < 0) {
        ESP_LOGE(TAG, "发送WebSocket帧头失败");
        return false;
    }
    
    // 发送掌码密钥
    uint8_t mask[4] = {0x12, 0x34, 0x56, 0x78};
    if (send(sock, mask, 4, 0) < 0) {
        ESP_LOGE(TAG, "发送WebSocket掌码失败");
        return false;
    }
    
    // 优化：直接在发送时进行掌码，避免额外的内存分配和拷贝
    size_t sent_total = 0;
    const size_t chunk_size = 1024;  // 分块发送，减少内存使用
    uint8_t chunk[chunk_size];
    
    while (sent_total < len) {
        size_t remaining = len - sent_total;
        size_t current_chunk = (remaining > chunk_size) ? chunk_size : remaining;
        
        // 对当前块进行掌码
        for (size_t i = 0; i < current_chunk; i++) {
            chunk[i] = data[sent_total + i] ^ mask[(sent_total + i) % 4];
        }
        
        // 发送当前块
        int sent = send(sock, chunk, current_chunk, 0);
        if (sent < 0) {
            ESP_LOGE(TAG, "发送WebSocket数据失败");
            return false;
        }
        
        sent_total += sent;
    }
    
    return true;
}


// WebSocket连接任务
static void websocket_task(void *pvParameters) {
    struct sockaddr_in server_addr;
    
    ESP_LOGI(TAG, "WebSocket连接任务开始");
    
    while (websocket_task_running) {
        // 创建socket
        websocket_sock = socket(AF_INET, SOCK_STREAM, 0);
        if (websocket_sock < 0) {
            ESP_LOGE(TAG, "创建WebSocket socket失败");
            vTaskDelay(5000 / portTICK_PERIOD_MS);
            continue;
        }

        // 显式关闭 TCP keepalive
        int ka = 0;
        setsockopt(websocket_sock, SOL_SOCKET, SO_KEEPALIVE, &ka, sizeof(ka));
        
        // 设置服务器地址
        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_port = htons(WEBSOCKET_SERVER_PORT);
        inet_pton(AF_INET, WEBSOCKET_SERVER_IP, &server_addr.sin_addr);
        
        // 连接到服务器
        ESP_LOGI(TAG, "连接到WebSocket服务器...");
        if (connect(websocket_sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
            ESP_LOGE(TAG, "WebSocket连接失败");
            close(websocket_sock);
            websocket_sock = -1;
            vTaskDelay(5000 / portTICK_PERIOD_MS);
            continue;
        }
        
        // WebSocket握手
        if (!websocket_handshake(websocket_sock)) {
            close(websocket_sock);
            websocket_sock = -1;
            vTaskDelay(5000 / portTICK_PERIOD_MS);
            continue;
        }
        
        // 等待一下让WebSocket连接稳定
        vTaskDelay(50 / portTICK_PERIOD_MS);
        
        websocket_connected = true;
        ESP_LOGI(TAG, "WebSocket连接成功，启动图片发送任务");
        
        // 启动图片帧发送任务
        start_websocket_frame_task();
        
        // 等待帧发送任务启动
        vTaskDelay(100 / portTICK_PERIOD_MS);
        
        // 保持连接活跃（简化版本，不使用PING/PONG）
        while (websocket_task_running && websocket_connected) {
            // 简单的连接状态检查
            vTaskDelay(1000 / portTICK_PERIOD_MS);
        }
        
        // 清理连接
        websocket_connected = false;
        stop_websocket_frame_task();
        
        // 使用互斥锁保护socket操作
        if (websocket_mutex && xSemaphoreTake(websocket_mutex, pdMS_TO_TICKS(1000)) == pdTRUE) {
            if (websocket_sock >= 0) {
                close(websocket_sock);
                websocket_sock = -1;
            }
            xSemaphoreGive(websocket_mutex);
        }
        
        ESP_LOGI(TAG, "WebSocket连接已断开，5秒后重连");
        vTaskDelay(5000 / portTICK_PERIOD_MS);
    }
    
    ESP_LOGI(TAG, "WebSocket连接任务退出");
    vTaskDelete(NULL);
}

// WebSocket图片帧发送任务
static void websocket_frame_task(void *pvParameters)
{
    camera_fb_t *fb = NULL;
    int frame_delay = 200;

    
    ESP_LOGI(TAG, "WebSocket图片帧发送任务开始运行");
    
    // 等待WebSocket连接稳定
    vTaskDelay(200 / portTICK_PERIOD_MS);
    
    while (websocket_frame_task_running && websocket_connected) {
        // 检查摄像头是否已初始化
        if (!is_camera_initialized()) {
            ESP_LOGW(TAG, "摄像头未初始化，等待初始化完成");
            vTaskDelay(1000 / portTICK_PERIOD_MS);
            continue;
        }
        
        // 检查socket状态
        bool socket_valid = false;
        if (websocket_mutex && xSemaphoreTake(websocket_mutex, pdMS_TO_TICKS(100)) == pdTRUE) {
            socket_valid = (websocket_sock >= 0);
            xSemaphoreGive(websocket_mutex);
        }
        
        if (!socket_valid) {
            ESP_LOGW(TAG, "WebSocket socket不可用，等待重连");
            vTaskDelay(1000 / portTICK_PERIOD_MS);
            continue;
        }
        
        fb = esp_camera_fb_get();
        if (fb) {
            uint8_t *jpg_buf = NULL;
            size_t jpg_len = 0;
            bool need_free = false;
            
            // 如果是RGB565格式，转换为JPEG
            if (fb->format == PIXFORMAT_RGB565) {
                if (frame2jpg(fb, 30, &jpg_buf, &jpg_len)) {
                    need_free = true;
                } else {
                    ESP_LOGE(TAG, "JPEG转换失败");
                    esp_camera_fb_return(fb);
                    continue;
                }
            } else {
                jpg_buf = fb->buf;
                jpg_len = fb->len;
            }
            
            // 通过WebSocket发送JPEG数据，使用互斥锁保护
            bool send_success = false;
            if (websocket_mutex && xSemaphoreTake(websocket_mutex, pdMS_TO_TICKS(50)) == pdTRUE) {
                if (websocket_sock >= 0 && websocket_connected) {
                    send_success = websocket_send_binary(websocket_sock, jpg_buf, jpg_len);
                }
                xSemaphoreGive(websocket_mutex);
            }
            
            if (send_success) {
                ESP_LOGI(TAG, "WebSocket发送图片成功，大小: %d bytes", jpg_len);
            } else {
                ESP_LOGE(TAG, "WebSocket发送图片失败");
                websocket_connected = false;
            }
            
            if (need_free && jpg_buf) {
                free(jpg_buf);
            }
            
            esp_camera_fb_return(fb);
        } else {
            ESP_LOGW(TAG, "获取摄像头帧失败");
            // 检查摄像头是否仍然有效
            sensor_t *s = esp_camera_sensor_get();
            if (s == NULL) {
                ESP_LOGE(TAG, "摄像头sensor获取失败，可能已断开连接");
                vTaskDelay(2000 / portTICK_PERIOD_MS);
                continue;
            }
            // 获取帧失败时短暂延时
            vTaskDelay(100 / portTICK_PERIOD_MS);
        }
        

        vTaskDelay(frame_delay / portTICK_PERIOD_MS);
    }
    
    ESP_LOGI(TAG, "WebSocket图片帧发送任务退出");
    websocket_frame_task_handle = NULL;
    vTaskDelete(NULL);
}

// 启动WebSocket图片帧发送任务
static void start_websocket_frame_task(void)
{
    if (!websocket_frame_task_running && websocket_frame_task_handle == NULL) {
        websocket_frame_task_running = true;
        BaseType_t result = xTaskCreatePinnedToCore(websocket_frame_task, "websocket_frame", 8192, NULL, 5, &websocket_frame_task_handle, 1);
        if (result == pdPASS) {
            ESP_LOGI(TAG, "WebSocket图片帧发送任务启动成功");
        } else {
            ESP_LOGE(TAG, "WebSocket图片帧发送任务启动失败");
            websocket_frame_task_running = false;
        }
    }
}

// 停止WebSocket图片帧发送任务
static void stop_websocket_frame_task(void)
{
    if (websocket_frame_task_running) {
        websocket_frame_task_running = false;
        // 等待任务自然结束
        int wait_count = 0;
        while (websocket_frame_task_handle != NULL && wait_count < 50) {
            vTaskDelay(100 / portTICK_PERIOD_MS);
            wait_count++;
        }
        ESP_LOGI(TAG, "WebSocket图片帧发送任务已停止");
    }
}

// 启动WebSocket客户端
void start_websocket_client(void)
{
    if (websocket_task_handle != NULL) {
        ESP_LOGW(TAG, "WebSocket客户端已在运行");
        return;
    }
    
    // 创建互斥锁
    if (websocket_mutex == NULL) {
        websocket_mutex = xSemaphoreCreateMutex();
        if (websocket_mutex == NULL) {
            ESP_LOGE(TAG, "创建WebSocket互斥锁失败");
            return;
        }
    }
    
    websocket_task_running = true;
    BaseType_t result = xTaskCreatePinnedToCore(websocket_task, "websocket_client", 8192, NULL, 5, &websocket_task_handle, 1);
    if (result == pdPASS) {
        ESP_LOGI(TAG, "WebSocket客户端任务启动成功");
    } else {
        ESP_LOGE(TAG, "WebSocket客户端任务启动失败");
        websocket_task_running = false;
    }
}

// 停止WebSocket客户端
void stop_websocket_client(void)
{
    if (websocket_task_running) {
        websocket_task_running = false;
        websocket_connected = false;
        
        // 使用互斥锁保护socket操作
        if (websocket_mutex && xSemaphoreTake(websocket_mutex, pdMS_TO_TICKS(1000)) == pdTRUE) {
            if (websocket_sock >= 0) {
                close(websocket_sock);
                websocket_sock = -1;
            }
            xSemaphoreGive(websocket_mutex);
        }
        
        // 等待任务结束
        int wait_count = 0;
        while (websocket_task_handle != NULL && wait_count < 50) {
            vTaskDelay(100 / portTICK_PERIOD_MS);
            wait_count++;
        }
        
        // 清理互斥锁
        if (websocket_mutex) {
            vSemaphoreDelete(websocket_mutex);
            websocket_mutex = NULL;
        }
        
        ESP_LOGI(TAG, "WebSocket客户端已停止");
    }
}
